# !/usr/bin/env python
# -*-coding:utf-8 -*-
"""
# File       : task_queue.py
# Time       ：2024/4/28 21:29
# Author     ：Y-aong
# version    ：python 3.7
# Description：
    redis任务队列
"""
import json
from collections import deque

import redis
from redis import Redis
from werkzeug.local import Local

from public.logger import logger


class LocalQueue:
    def __init__(self):
        self._local = Local()

    def add(self, queue_name: str, fields: dict):
        """将任务加入队列中"""
        task_queue = getattr(self._local, queue_name, None)
        if task_queue is None:
            task_queue = deque()
            setattr(self._local, queue_name, task_queue)
        task_queue.append(fields)
        return task_queue

    def get(self, queue_name):
        """获取队列中的数据"""
        task_queue = getattr(self._local, queue_name, None)
        if not task_queue:
            return None
        return task_queue.popleft()

    def is_empty(self, queue_name) -> bool:
        task_queue = getattr(self._local, queue_name, None)
        return task_queue is None or (isinstance(task_queue, deque) and len(task_queue) == 0)

    def ack(self, queue_name, task_id: str):
        task_queue: deque = getattr(self._local, queue_name, None)
        if not task_queue:
            return
        for item in task_queue:
            if item == task_id:
                task_queue.remove(item)


class RedisQueue:
    def __init__(self, host: str = '127.0.0.1', port: int = 6379, db=0, username=None, password=None):
        if username and password:
            self.redis = Redis(host=host, port=port, username=username, password=password, db=db, decode_responses=True)
        else:
            self.redis = Redis(host=host, port=port, db=db, decode_responses=True)

        self.fields = set()

    def _check_message(self, fields: dict):
        message = {}
        for key, val in fields.items():
            if isinstance(val, (list, dict)) or val is None:
                val = json.dumps(val)
                self.fields.add(key)
            message[key] = val
        return message

    def _handle_message(self, message):
        response_message = {}

        for key, val in message.items():
            if key in self.fields:
                val = json.loads(val)
            response_message[key] = val
        return response_message

    def add(self, stream_name: str, message: dict, task_id='*'):
        """
        消息队列添加数据
        @param stream_name:
        @param message:
        @param task_id:
        @return:
        """
        message = self._check_message(message)
        try:
            response = self.redis.xadd(stream_name, message, id=task_id)
        except redis.exceptions.RedisError as e:
            logger.error(f"Failed to add message to stream: {e}")
            response = None
        return response

    def create_group(self, stream_name: str, consumer_group: str, stream_id='0', mk_stream=True):
        """创建消费组"""
        try:
            response = self.redis.xgroup_create(stream_name, consumer_group, id=stream_id, mkstream=mk_stream)
        except redis.exceptions.RedisError as e:
            logger.error(f"Failed to create consumer group: {e}")
            response = None
        return response

    def group_read(self, stream_name, consumer_name: str, group_name: str, count: int = 1, block: int = 5000):
        """
        从消费组读取消息
        @param stream_name:
        @param consumer_name:
        @param group_name:
        @param count:
        @param block:
        @return:
        """
        try:
            response = self.redis.xreadgroup(
                groupname=group_name, consumername=consumer_name, streams={stream_name: '>'}, count=count, block=block
            )
            messages = []
            for stream, message_data in response:
                for message_id, message in message_data:
                    messages.append({
                        'message_id': message_id,
                        'message': self._handle_message(message)
                    })
        except redis.exceptions.RedisError as e:
            logger.error(f"Failed to read messages from group: {e}")
            messages = []
        return messages

    def ack(self, stream_name: str, group_name: str, message_id: str, delete=True):
        """
        确认消息已被消费
        @param stream_name:
        @param group_name:
        @param message_id:
        @param delete:是否删除掉当前消息
        @return:
        """
        if not message_id:
            logger.warn(f'Current message {message_id} has not been acknowledged')
            return
        try:
            self.redis.xack(stream_name, group_name, message_id)
            if delete:
                self.delete_message(stream_name, message_id)
        except redis.exceptions.RedisError as e:
            logger.error(f"Failed to acknowledge message: {e}")

    def get_stream_info(self, stream_name: str):
        """
        获取Stream的信息
        @param stream_name:
        @return:
        """
        try:
            return self.redis.xinfo_stream(stream_name)
        except redis.exceptions.RedisError as e:
            logger.error(f"Failed to get stream info: {e}")
            return None

    def delete_message(self, stream_name: str, message_id: str):
        """
        删除消息
        @param stream_name:
        @param message_id:
        @return:
        """
        try:
            return self.redis.xdel(stream_name, message_id)
        except redis.exceptions.RedisError as e:
            logger.error(f"Failed to delete message: {e}")
            return None

    def get_pending_info(self, stream_name: str, group_name: str):
        """
        获取待处理消息的信息
        @param stream_name:
        @param group_name:
        @return:
        """
        try:
            return self.redis.xpending(stream_name, group_name)
        except redis.exceptions.RedisError as e:
            logger.error(f"Failed to get pending info: {e}")
            return {}

    def get_group_info(self, stream_name):
        """
        获取消费者组的信息
        @param stream_name:
        @return:
        """
        try:
            return self.redis.xinfo_groups(stream_name)
        except redis.exceptions.RedisError as e:
            logger.error(f"Failed to get group info: {e}")
            return []

    def is_empty(self, stream_name: str):
        """判断队列是否为空"""
        try:
            return self.redis.xlen(stream_name) == 0
        except redis.exceptions.RedisError as e:
            logger.error(f"Failed to get pending info: {e}")
            return False
